from collections import deque
from pythongo.ui import BaseStrategy
from pythongo.base import BaseParams, BaseState
from pythongo.core import KLineStyleType
from pythongo.utils import KLineGenerator
from pythongo.indicator import Indicators
from pydantic import Field
import talib
import numpy as np

# ---------- 参数 ----------
class Params(BaseParams):
    exchange: str = Field(default="SHFE", title="交易所代码")
    instrument_id: str = Field(default="IF2309", title="合约代码")
    kline_style: KLineStyleType = Field(default="M1", title="K线周期")
    boll_period: int = Field(default=20, title="布林线周期", ge=1)
    boll_std_multiplier: int = Field(default=2, title="布林线标准差倍数", ge=1)
    macd_fast_period: int = Field(default=12, title="MACD快线周期", ge=1)
    macd_slow_period: int = Field(default=26, title="MACD慢线周期", ge=1)
    macd_signal_period: int = Field(default=9, title="MACD信号周期", ge=1)
    order_volume: int = Field(default=1, title="每次开仓手数", ge=1)
    capital_threshold: float = Field(default=50.0, title="资金监控阈值（百分比）")

class State(BaseState):
    last_macd: float = Field(default=0.0, title="上一个MACD值")
    has_position: bool = Field(default=False, title="是否有仓位")
    position: int = Field(default=0, title="当前仓位")
    stop_loss_price: float = Field(default=0.0, title="止损价格")
    take_profit_price: float = Field(default=0.0, title="止盈价格")

class BollReversal_v2(BaseStrategy):
    def __init__(self):
        super().__init__()
        self.params_map = Params()
        self.state_map = State()
        self.kline_generator: KLineGenerator = None
        self.close_buffer = deque(maxlen=self.params_map.boll_period)
        self.trading_paused = False
        self.initial_equity = None  # 初始化为 None

    def on_start(self):
        super().on_start()

        investor = self.get_investor_data(1).investor_id
        account = self.get_account_fund_data(investor)
        self.initial_equity = account.balance + account.close_profit  # 记录初始权益
        self.output(f"初始权益: {self.initial_equity}")
        
        self.kline_generator = KLineGenerator(
            exchange=self.params_map.exchange,
            instrument_id=self.params_map.instrument_id,
            style=self.params_map.kline_style,
            callback=self.on_new_kline,
            real_time_callback=self.on_realtime_kline
        )
        self.kline_generator.push_history_data()
        self.output("策略启动，开始接收K线...")

    def on_stop(self):
        super().on_stop()
        self.output("策略停止")

    def on_new_kline(self, kline):
        self.close_buffer.append(kline.close)
        if len(self.close_buffer) < self.params_map.boll_period:
            return
        self.calculate_and_plot(kline)

    def on_realtime_kline(self, kline):
        self.on_new_kline(kline)

    def calculate_and_plot(self, kline):
        if self.state_map.has_position:
            return

        closes = list(self.close_buffer)
        arr = np.array(closes, dtype=float)

        middle = talib.SMA(arr, self.params_map.boll_period)
        std = talib.STDDEV(arr, self.params_map.boll_period)
        upper = float(middle[-1] + std[-1] * self.params_map.boll_std_multiplier)
        lower = float(middle[-1] - std[-1] * self.params_map.boll_std_multiplier)

        macd, signal, hist = talib.MACD(
            arr,
            fastperiod=self.params_map.macd_fast_period,
            slowperiod=self.params_map.macd_slow_period,
            signalperiod=self.params_map.macd_signal_period
        )

        self.widget.recv_kline({
            "kline": kline,
            "upper": round(upper, 2),
            "middle": round(float(middle[-1]), 2),
            "lower": round(lower, 2),
            "macd": round(float(macd[-1]), 4),
            "signal": round(float(signal[-1]), 4),
            "hist": round(float(hist[-1]), 4)
        })

        self.state_map.last_macd = float(hist[-1])
        self.update_status_bar()

        self.check_global_risk(kline)
        self.check_signals(kline, lower, middle[-1], upper, macd[-1])

    def check_global_risk(self, kline):
        if self.trading_paused:
            return

        investor = self.get_investor_data(1).investor_id
        account = self.get_account_fund_data(investor)

        current_equity = account.balance + account.close_profit
        used_margin = account.margin  # 使用正确的属性名
        total_equity = account.balance + account.close_profit

        if self.initial_equity is None:
            self.output("初始权益未初始化，无法计算单日亏损")
            return

        daily_loss = (self.initial_equity - current_equity) / self.initial_equity
        usage_ratio = used_margin / total_equity if total_equity > 0 else 0

        threshold = self.params_map.capital_threshold / 100.0

        if daily_loss >= threshold:
            self.output(f"单日亏损熔断: {daily_loss:.2%}")
            self.emergency_close_all("单日亏损熔断")
            return

        if usage_ratio >= threshold:
            self.output(f"资金使用率熔断: {usage_ratio:.2%}")
            self.trading_paused = True

    def check_signals(self, kline, lower, middle, upper, macd):
        if self.state_map.has_position:
            return

        if kline.close > middle and kline.close < upper:
            if macd < 0 or (macd == 0 and self.state_map.last_macd < 0):
                self.open_position("SELL", kline, upper, lower)
        elif kline.close < middle and kline.close > lower:
            if macd > 0 or (macd == 0 and self.state_map.last_macd > 0):
                self.open_position("BUY", kline, lower, upper)

    def open_position(self, direction, kline, stop_loss, take_profit):
        order_direction = "BUY" if direction == "BUY" else "SELL"
        order_id = self.send_order(
            exchange=self.params_map.exchange,
            instrument_id=self.params_map.instrument_id,
            volume=self.params_map.order_volume,
            price=0,  # 市价
            order_direction=order_direction,
            market=True
        )

        if order_id != -1:
            self.state_map.has_position = True
            self.state_map.position = self.params_map.order_volume if direction == "BUY" else -self.params_map.order_volume
            self.state_map.stop_loss_price = round(stop_loss, 2)
            self.state_map.take_profit_price = round(take_profit, 2)
            self.output(f"开仓成功 - 方向:{direction}, 止损:{stop_loss}, 止盈:{take_profit}")

    def on_tick(self, tick):
        if self.kline_generator:
            self.kline_generator.tick_to_kline(tick)
        self.check_exit_conditions(tick)

    def check_exit_conditions(self, tick):
        if not self.state_map.has_position:
            return

        price = tick.last_price
        direction = "LONG" if self.state_map.position > 0 else "SHORT"
        stop_loss = self.state_map.stop_loss_price
        take_profit = self.state_map.take_profit_price

        if direction == "LONG":
            if price <= stop_loss:
                self.close_position("多头止损")
            elif price >= take_profit:
                self.close_position("多头止盈")
        else:
            if price >= stop_loss:
                self.close_position("空头止损")
            elif price <= take_profit:
                self.close_position("空头止盈")

    def close_position(self, reason):
        order_direction = "SELL" if self.state_map.position > 0 else "BUY"
        order_id = self.auto_close_position(
            exchange=self.params_map.exchange,
            instrument_id=self.params_map.instrument_id,
            volume=abs(self.state_map.position),
            price=0,
            order_direction=order_direction,
            market=True
        )

        if order_id != -1:
            self.state_map.has_position = False
            self.state_map.position = 0
            self.state_map.stop_loss_price = 0.0
            self.state_map.take_profit_price = 0.0
            self.output(f"{reason} - 平仓成功")

    def emergency_close_all(self, reason):
        self.close_position(reason)
        self.trading_paused = True
        self.pause_strategy()